#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将多个 TSV 文件写入同一个 Excel 文件中。

- 第一个 TSV 写入 Sheet1
- 空一行后写入第二个 TSV（同样位于 Sheet1）
- 第三个 TSV 写入 Sheet2
- 第四个 TSV 写入 Sheet3

读取 TSV 时严格把 "NA" 当作普通字符串，而非缺失值。
"""

import pandas as pd
from openpyxl import load_workbook
from openpyxl.utils.dataframe import dataframe_to_rows
from pathlib import Path
from argparse import ArgumentParser, Namespace


def parse_args() -> Namespace:
    """解析命令行参数"""
    parser = ArgumentParser(description="将多个 TSV 写入 Excel 文件")

    parser.add_argument("--tsv1", required=True, help="第一个 TSV 文件（写入 Sheet1）")
    parser.add_argument("--tsv2", required=True, help="第二个 TSV 文件（写入 Sheet1，空一行后追加）")
    parser.add_argument("--tsv3", required=True, help="第三个 TSV 文件（写入 Sheet2）")
    parser.add_argument("--tsv4", required=True, help="第四个 TSV 文件（写入 Sheet3）")
    parser.add_argument("--out_xlsx", required=True, help="输出 Excel 文件路径")

    return parser.parse_args()


def write_single_sheet(ws, df, start_row=1, header=True):
    """
    将 DataFrame 写入指定 sheet，从指定行开始。

    Parameters
    ----------
    ws : openpyxl worksheet
        要写入的工作表
    df : pandas.DataFrame
        数据内容
    start_row : int
        写入的起始行
    header : bool
        是否写入表头
    """
    rows = dataframe_to_rows(df, index=False, header=header)
    row_i = start_row

    for row in rows:
        for col_j, cell_value in enumerate(row, start=1):
            ws.cell(row=row_i, column=col_j, value=cell_value)
        row_i += 1

    return row_i  # 返回写入结束的下一行位置


def write_tsv_to_excel(tsv1, tsv2, tsv3, tsv4, out_xlsx):
    """
    读取 4 个 TSV 文件并按要求写入 Excel。
    """
    out_xlsx = Path(out_xlsx).absolute()

    # 读取 TSV，强制所有内容当作字符串（"NA" 不转为 NaN）
    df1 = pd.read_csv(tsv1, sep="\t", dtype=str, na_filter=False, header=None)
    df2 = pd.read_csv(tsv2, sep="\t", dtype=str, na_filter=False)
    df3 = pd.read_csv(tsv3, sep="\t", dtype=str, na_filter=False)
    df4 = pd.read_csv(tsv4, sep="\t", dtype=str, na_filter=False)

    # Step 1：先创建 Excel 并写入第一个 TSV
    with pd.ExcelWriter(out_xlsx, engine="openpyxl") as writer:
        df1.to_excel(writer, index=False, sheet_name="Sheet1", header=False)

    # Step 2：打开 Excel，手工写入第二个 TSV
    wb = load_workbook(out_xlsx)
    ws1 = wb["Sheet1"]

    # 找到最后一行，为第二个 TSV 空出一行
    start_row = ws1.max_row + 2

    # 写第二个 TSV
    write_single_sheet(ws1, df2, start_row=start_row, header=True)

    # Step 3：第三个 TSV 写入 Sheet2
    ws2 = wb.create_sheet("Sheet2")
    write_single_sheet(ws2, df3, start_row=1, header=True)

    # Step 4：第四个 TSV 写入 Sheet3
    ws3 = wb.create_sheet("Sheet3")
    write_single_sheet(ws3, df4, start_row=1, header=True)

    wb.save(out_xlsx)


def main():
    args = parse_args()
    Path(args.out_xlsx).parent.mkdir(parents=True, exist_ok=True)
    write_tsv_to_excel(
        args.tsv1,
        args.tsv2,
        args.tsv3,
        args.tsv4,
        args.out_xlsx
    )


if __name__ == "__main__":
    main()
